package marek.ai.rl;

import java.util.ArrayList;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Random;

public class ReinforcedLearning {
	public static final double LEARN_RATE = 0.9;
	public static class State {
		public final static int POSSIBLE_POSITIONS = 3;
		private int gY;
		private int gX;

		public void setgY(int gY) {
			this.gY = gY;
		}

		public int getgY() {
			return gY;
		}

		public void setgX(int gX) {
			this.gX = gX;
		}

		public int getgX() {
			return gX;
		}

		@Override
		public boolean equals(Object o) {
			if (!(o instanceof State)) {
				return false;
			}
			State s = (State) o;
			return (s.gX == this.gX && s.gY == this.gY);
		}

		@Override
		public int hashCode() {
			return (gX+1)*10 + (gY+1);
		}
		@Override
		public String toString () {
			return "[" + gX + ", " + gY + "]";
		}

	}

	public static class Action {
		public static class Movement {
			private int xOffset;
			private int yOffset;
			
			private Movement(int xOffset, int yOffset) {
				this.xOffset = xOffset;
				this.yOffset = yOffset;
			}
			
			private State executeOn(State s) {
				State newState = new State();
				newState.setgX(s.getgX() + xOffset);
				newState.setgY(s.getgY() + yOffset);
				if (newState.getgX() < 0 || newState.getgX() >= State.POSSIBLE_POSITIONS || newState.getgY() < 0 || newState.getgY() >= State.POSSIBLE_POSITIONS) {
					throw new IllegalStateException(newState.toString());
				}
				return newState;
			}
		}
		
		public final static Movement UP = new Movement(0,-1);
		public final static Movement DOWN = new Movement(0,1);
		public final static Movement RIGHT = new Movement(1,0);
		public final static Movement LEFT = new Movement(-1,0);
		public final static Movement[] MOVEMENTS = {UP, DOWN, RIGHT, LEFT};
		
		private Movement actionMovement;
		
		public Action(Movement m) {
			this.actionMovement = m;
		}

		public State doAction(State s) {
			return actionMovement.executeOn(s);
		}

		@Override
		public boolean equals(Object o) {
			if (!(o instanceof Action)) {
				return false;
			}
			Action a = (Action) o;
			return (a.actionMovement.xOffset == this.actionMovement.xOffset && a.actionMovement.yOffset == this.actionMovement.yOffset);
		}

		@Override
		public int hashCode() {
			return (actionMovement.xOffset+2)*10 + (actionMovement.yOffset+2);
		}
	}

	public static class StateSpace {
		private Map<State, List<Action>> stateToActions;
		private List<State> stateList;
		private int nb2Actions = 0;
		private int nb3Actions = 0;
		private int nb4Actions = 0;

		public StateSpace() {
			stateToActions = new HashMap<State, List<Action>>();
			stateList = new ArrayList<State>();
			for (int y = 0; y < State.POSSIBLE_POSITIONS; y++) {
				for (int x = 0; x < State.POSSIBLE_POSITIONS; x++) {
					State s = new State();
					s.setgX(x);
					s.setgY(y);
					stateList.add(s);

					List<Action> possibleActionForS = new ArrayList<Action>();
					stateToActions.put(s, possibleActionForS);

					// Try move all directions from current state s
					for (Action.Movement m : Action.MOVEMENTS) {
						try {
							Action a = new Action(m);
							a.doAction(s);
							// if no InvalidState add a to possible actions for s
							possibleActionForS.add(a);
						} catch (IllegalStateException e) {
							// Ignore invalid transitions
						}
					}
					switch (possibleActionForS.size()) {
					case 2:
						nb2Actions++;
						break;
					case 3:
						nb3Actions++;
						break;
					case 4:
						nb4Actions++;
						break;
					}
				}
			}
		}

		public List<Action> getPossibleActions(State s) {
			return stateToActions.get(s);
		}

		public int getNbPolicies() {
			int total = (int) Math.pow(2, nb2Actions);
			if (nb3Actions > 0)
				total *=(int) Math.pow(3, nb3Actions);
			if (nb4Actions > 0)
				total *=(int) Math.pow(4, nb4Actions);
			return total;
		}

		public int getNbStates() {
			return stateList.size();
		}
				
		public List<State> getStates() {
			return stateList;
		}

		public State selectRandomState() {
			Random r = new Random(System.currentTimeMillis());
			int stateIndex = r.nextInt(getNbStates());
			return stateList.get(stateIndex);
		}
	}


	/* Dynamic Value Iteration method
	 * 
	 */
	public static class V {
		/*
		 * Policy should only be instantiated by the compile method of V
		 */
		public static class Policy {
			private Map<State, Action[]> stateActions;
			private Policy() {
				
			}
		}

		private Map<State, Double> stateValue = new HashMap<State, Double>();
		public double getValue(State s) {
			Double i = stateValue.get(s);
			if (i == null) {
				return 0;
			}
			return i;
		}
		private void setValue (State s, double value) {
			stateValue.put(s, value);
		}
		public Policy compile() {
			// Iterate over states and their values to generate policy based on V(s)
			return null;
		}
		
		@Override
		public String toString() {
			String s = "";
			for (State state : stateValue.keySet()) {
				s += state.toString() + " = " + getValue(state) + "\r\n";
			}
			return s;
		}
		
		public V clone() {
			V v = new V();
			for (State s : this.stateValue.keySet()) {
				v.setValue(s, this.getValue(s));
			}
			return v;
		}
		
		@Override
		public boolean equals(Object o) {
			if (!(o instanceof V)) {
				return false;
			}
			V v = (V)o;
			if (v.stateValue.keySet().size() !=  this.stateValue.keySet().size()) {
				return false;
			}
			for (State s : this.stateValue.keySet()) {
				if (v.getValue(s) != this.getValue(s)) {
					return false;
				}
			}
			return true;
		}
	}

	public static void valueIteration(StateSpace stateSpace, V v) {
		// loop while v(s) does not change anymore
		V prevV = null;
		while (prevV == null |! v.equals(prevV)) {
			prevV = v.clone();
			for (State state : stateSpace.getStates()) {
				List<Action> possibleActions = stateSpace.getPossibleActions(state);
				double max = -1;
				for (Action action : possibleActions) {
					ActionResult ar = executeAction(state, action);
					// new v(s) value
					double newPotentialValue = ar.reward + LEARN_RATE * v.getValue(ar.newState);
					if (newPotentialValue > max) {
						max = newPotentialValue;
					}
				}
				v.setValue(state, max);
			}
		}
	}
	
	public static class ActionResult {
		private int reward;
		private State newState;
	}

	/*
	 * Execute action : get new state and reward
	 */
	public static ActionResult executeAction(State currentState, Action action) {
		State newState = action.doAction(currentState);
		int reward = 0;
		// if vertical joined arm is down (touching ground)
		if (currentState.getgY() == State.POSSIBLE_POSITIONS-1 && newState.getgY() == State.POSSIBLE_POSITIONS-1){
			// reward =  diff on x from current / to new state
			reward = currentState.getgX() - newState.getgX();
		}  // else leave reward = 0;
		ActionResult ar = new ActionResult();
		ar.reward = reward;
		ar.newState = newState;
		return ar;
	}

	public static void main(String[] args) {
		StateSpace sp = new StateSpace();
		V v = new V();
		valueIteration(sp, v);
		
		System.out.println("Nb states: " + sp.getNbStates() + " - Nb policies: " + sp.getNbPolicies());
		System.out.println("V(allStates):\r\n " + v);
	}
}
